package org.examples.java.util.concurrent;

import java.util.ArrayList;
import java.util.Collection;
import java.util.LinkedList;
import java.util.List;
import java.util.Random;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.RecursiveTask;

import org.junit.Test;

public class ForkJoinTest {

    private static final int WIDTH = 2000;//10;
    private static final int DEPTH = 2;//8;
    private static final int CREATOR_DEPTH = DEPTH - 1;
    long start = System.currentTimeMillis();

    //@Test
    public void testTree() {
        for (int d = 0; d <= DEPTH; d++) {
            for (int w = 0; w <= WIDTH; w++) {
                createTree(d, w);
                System.out.println("Tree created depth=" + d + ", width=" + w +
                        "\t\t Time: " + elapsedTime() + "ms");
            }
        }
    }

    @Test
    public void testCompute() {
        Node root = createTree(DEPTH, WIDTH);
        System.out.println("Tree created (seq) depth=" + DEPTH + ", width=" + WIDTH + ". Time: " + elapsedTime() + "ms");
        root = new ForkJoinPool().invoke(new Creator(DEPTH, WIDTH));
        System.out.println("Tree created (f/j) depth=" + DEPTH + ", width=" + WIDTH + ". Time: " + elapsedTime() + "ms");

        for (int i = 0; i < 5; i++) {
            Long count = new ForkJoinPool().invoke(new Counter(root, 1));
            System.out.println("Total nodes (f/j): " + count + ". Time: " + elapsedTime() + "ms");
            System.out.println("Total nodes (seq): " + root.countAll() + ". Time: " + elapsedTime() + "ms");
        }
    }

    private static Node createTree(int depth, int width) {
        if (depth == 0)
            return null;
        Node node = new Node();

        int w = width;// Math.max(1, new Random().nextInt(width) + 1);
        for (int i = 0; i < w; i++)
            node.addChildren(createTree(depth - 1, width));

        return node;
    }

    static class Node {
        private Collection<Node> children = new ArrayList<>();
        private long value = new Random().nextInt(1000);

        Collection<Node> getChildren() {
            return children;
        }

        public long countAll() {
            long count = 1;
            for (Node node : children)
                count += node.countAll();
            return count;
        }

        public void addChildren(Node node) {
            if (node != null)
                children.add(node);
        }

        long getValue() {
            return value;
        }

        @Override
        public String toString() {
            return toString(0);
        }

        private String toString(int depth) {
            StringBuilder str = new StringBuilder();
            for (int i = 0; i < depth; i++)
                str.append(".");
            str.append(value);
            str.append("\n");
            for (Node node : children)
                str.append(node.toString(depth + 1));
            return str.toString();
        }
    }

    @SuppressWarnings("serial")
    static class Counter extends RecursiveTask<Long> {
        private final Node node;
        private int maxDepth;

        Counter(Node node, int maxDepth) {
            this.node = node;
            this.maxDepth = maxDepth;
        }

        @Override
        protected Long compute() {
            if (maxDepth == 0) {
                return node.countAll();
            }

            long count = 1;
            List<Counter> subTasks = new LinkedList<>();

            for (Node child : node.getChildren()) {
                Counter task = new Counter(child, maxDepth - 1);
                task.fork();
                subTasks.add(task);
            }

            for (Counter task : subTasks)
                count += task.join();

            return count;
        }

    }

    @SuppressWarnings("serial")
    static class Creator extends RecursiveTask<Node> {
        private final int depth, width;

        Creator(int depth, int width) {
            this.depth = depth;
            this.width = width;
        }

        @Override
        protected Node compute() {
            if (depth == CREATOR_DEPTH)
                return createTree(depth, width);

            Node node = new Node();
            List<Creator> subTasks = new LinkedList<>();
            for (int i = 0; i < width; i++) {
                Creator task = new Creator(depth - 1, width);
                task.fork();
                subTasks.add(task);
            }

            for (Creator task : subTasks)
                node.addChildren(task.join());

            return node;
        }
    }

    private long elapsedTime() {
        long elapsedTime = System.currentTimeMillis() - start;
        start = System.currentTimeMillis();
        return elapsedTime;
    }
}
